import os
import random
import shutil

# 设置路径
image_dir = 'task/train'
label_dir = 'task/label'
base_dir = 'task'
train_image_dir = os.path.join(base_dir, 'train_set', 'images')
train_label_dir = os.path.join(base_dir, 'train_set', 'labels')
test_image_dir = os.path.join(base_dir, 'test_set', 'images')
test_label_dir = os.path.join(base_dir, 'test_set', 'labels')

# 创建目录结构
os.makedirs(train_image_dir, exist_ok=True)
os.makedirs(train_label_dir, exist_ok=True)
os.makedirs(test_image_dir, exist_ok=True)
os.makedirs(test_label_dir, exist_ok=True)

# 获取所有标注文件
label_files = [f for f in os.listdir(label_dir) if f.endswith('.txt')]
randon.seed(202409260926)
random.shuffle(label_files)  # 随机打乱标注文件顺序

# 计算划分索引
split_index = int(len(label_files) * 0.8)  # 80% 用于训练

# 划分训练集和测试集
train_labels = label_files[:split_index]
test_labels = label_files[split_index:]

# 复制训练集文件
for label in train_labels:
    base_name = os.path.splitext(label)[0]
    image_file = os.path.join(image_dir, base_name + '.jpg')
    
    # 检查对应的图像文件是否存在
    if os.path.exists(image_file):
        shutil.copy(image_file, train_image_dir)
        shutil.copy(os.path.join(label_dir, label), train_label_dir)
    else:
        print(f"跳过标注文件 {label}：对应的图像文件不存在。")

# 复制测试集文件
for label in test_labels:
    base_name = os.path.splitext(label)[0]
    image_file = os.path.join(image_dir, base_name + '.jpg')
    
    # 检查对应的图像文件是否存在
    if os.path.exists(image_file):
        shutil.copy(image_file, test_image_dir)
        shutil.copy(os.path.join(label_dir, label), test_label_dir)
    else:
        print(f"跳过标注文件 {label}：对应的图像文件不存在。")

print("训练集和测试集划分完成。")